import pdfplumber
import pandas as pd
import os

if __name__ == "__main__":

    # 定义要提取的营养成分列表
    nutrients_to_extract = [
        "能量(kJ)",
        "蛋白质(g)",
        "脂肪(g)",
        "碳水化合物(g)",
        "钠(mg)",
        "氯(mg)",
        "钾(mg)",
        "磷(mg)",
    ]

    # 初始化一个空列表，用于存储每个文件的提取结果
    results = []

    # 假设你的PDF文件都在这个文件夹下，你可以根据实际情况修改路径
    folder_path = "/Users/lhmily/Desktop/2024data/data/Food_Manual"
    # folder_path = "/Users/bytedance/Downloads/2024data/data/Food_Manual"
    pdf_files = [f for f in os.listdir(folder_path) if f.endswith(".pdf")]

    # 检查保存结果的文件夹是否存在，不存在则创建
    save_folder = os.path.dirname(os.path.join(folder_path, "result1.xlsx"))
    if not os.path.exists(save_folder):
        os.makedirs(save_folder)

    for pdf_file in pdf_files:
        file_path = os.path.join(folder_path, pdf_file)
        registration_number = os.path.splitext(pdf_file)[0]
        data = {"注册证号": registration_number}
        with pdfplumber.open(file_path) as pdf:
            for page in pdf.pages:
                tables = page.extract_tables()
                for table in tables:
                    if len(table) >= 2:
                        df = pd.DataFrame(table[1:], columns=table[0])
                        if "营养成分" in df.columns and "每100kJ" in df.columns:

                            for nutrient in nutrients_to_extract:
                                # selected = df.loc[df['营养成分'] == nutrient, '每100kj']
                                target_row = df[df['营养成分'] == nutrient]
                                if not target_row.empty:
                                    target_column = target_row['每100kJ']
                                    if len(target_column) > 0:
                                        target_value = target_column.values[0]
                                        data[nutrient] = target_value                     
        results.append(data)

    # 将结果转换为DataFrame
    df_results = pd.DataFrame(results)

    # 找出每100kJ中蛋白质含量最高的三种特医食品
    top_3_protein = df_results.sort_values(by="蛋白质(g)", ascending=False).head(3)

    # 保存结果到Excel文件
    excel_file_path = os.path.join(folder_path, "result1.xlsx")
    with pd.ExcelWriter(excel_file_path) as writer:
        df_results.to_excel(writer, sheet_name="All Data", index=False)